# Author: Xinwu
# Describe: 
# Completion Time: 
# Email: lexinwu@outlook.com

if (!require("optparse", quietly = TRUE))
  install.packages("optparse", repo = "https://mirrors.tuna.tsinghua.edu.cn/CRAN/")

usage = "Rscript inferCNV.R --rdata project.rdata --normal_cell_clusters 1,3,5 --cancer_cell_clusters 2,4,6,7,8,9 --cluster_type seurat_clusters --gene_pos_file mouse_gene_pos.bed"
option_list = list(
  make_option("--rdata",type = "character",action = "store",default = NULL,help = "project.rdata 文件"),
  make_option("--normal_cell_clusters",type = "character",action = "store",default = NULL,help = "正常细胞所属的 Cluster, 逗号分隔"),
  make_option("--cancer_cell_clusters",type = "character",action = "store",default = NULL,help = "肿瘤细胞所属的 Cluster, 逗号分隔"),
  make_option("--cluster_type",type = "character",action = "store",default = "CellType",help = "cluster 注释的类型, 一般有 seurat_clusters, Cluster, Celltype"),
  make_option("--gene_pos_file",type = "character",action = "store",default = NULL,help = "基因的位置文件"),
  make_option("--thread",type = "integer",action = "store",default = 20,help = "thread number default: 20")
)
opt_parser = OptionParser(option_list = option_list,usage = usage)
opt = parse_args(opt_parser)

sava_pdf <- function(obj = NULL,file = NULL,
                     width = 6, height = 5){
  pdf(file = file, width = width, height = height)
  print(obj)
  dev.off()
  file_png <- gsub(pattern = "pdf",replacement = "png",x = file)
  png(file = file_png, width = width, height = height, res = 300, units = "in")
  print(obj)
  dev.off()
}

create_obj_rdata <- function(project = NULL, gene_pos_file = NULL,
                             normal_cell = NULL, cancer_cell = NULL,
                             cluster_type = "CellType"){
  library(Seurat)
  Idents(project) <- cluster_type
  project <- subset(project, idents = c(normal_cell,cancer_cell))
  raw_count <- project@assays$RNA@counts
  raw_count <- raw_count[Matrix::rowSums(raw_count) > 0,]
  gene_pos <- data.table::fread(gene_pos_file, sep = "\t", header = TRUE)
  gene <- gene_pos$geneName
  comm_genes <- intersect(rownames(raw_count), gene)
  raw_count <- raw_count[comm_genes, ]
  gene_pos <- gene_pos[as.vector(gene_pos$geneName) %in% comm_genes, ]
  gene_pos <- gene_pos[!duplicated(gene_pos$geneName),]
  gene_pos <- tibble::column_to_rownames(gene_pos, var = "geneName")
  normal_cell_indices <- list()
  for(a_normal in normal_cell){
    if(a_normal %in% unique(as.vector(project[[cluster_type]][,1]))){
      normal_cell_indices[[a_normal]] <- which(as.vector(project[[cluster_type]][,1]) == a_normal)
    }
  }
  cancer_cell_indices <- list()
  for(a_tumor in cancer_cell){
    if(a_tumor %in% unique(as.vector(project[[cluster_type]][,1]))){
      cancer_cell_indices[[a_tumor]] <- which(as.vector(project[[cluster_type]][,1]) == a_tumor)
    }
  }
  names(gene_pos) <- c(infercnv:::C_CHR, infercnv:::C_START, infercnv:::C_STOP)
  raw_count <- as.matrix(raw_count)
  order_ret <- infercnv:::.order_reduce(data = raw_count, genomic_position = gene_pos)
  raw_data <- order_ret$expr
  input_gene_order <- order_ret$order
  infercnv_obj <- new(Class = "infercnv", 
                      expr.data = raw_data, 
                      count.data = raw_data, 
                      gene_order = input_gene_order, 
                      reference_grouped_cell_indices = normal_cell_indices, 
                      observation_grouped_cell_indices = cancer_cell_indices,
                      tumor_subclusters = NULL, .hspike = NULL)
  return(infercnv_obj)
}

run_cnv <- function(obj = NULL, thread = 20){
  infercnv_obj <- infercnv::run(obj, cutoff = 0.1, out_dir = "./", 
                                cluster_by_groups = T, denoise = T, HMM = TRUE, 
                                num_threads = thread, no_prelim_plot = TRUE, no_plot = TRUE, 
                                BayesMaxPNormal = -1, analysis_mode = 'subclusters', 
                                tumor_subcluster_partition_method = 'random_trees',
                                tumor_subcluster_pval = 0.05)
  infercnv::plot_cnv(infercnv_obj, k_obs_groups = 1, cluster_by_groups = TRUE,
                     cluster_references = TRUE, out_dir = './', title = "inferCNV",
                     output_filename = "infercnv", output_format = 'pdf',
                     write_expr_matrix = TRUE, useRaster = TRUE)
  
}

library(infercnv)
load(opt$rdata)
NormalCellType <- opt$normal_cell_clusters
NormalCellType <- strsplit(NormalCellType, ',')[[1]]
TumorCellType <- opt$cancer_cell_clusters
TumorCellType <- strsplit(TumorCellType, ',')[[1]]
obj <- create_obj_rdata(project = project_Sob, gene_pos_file = opt$gene_pos_file,
                        normal_cell = NormalCellType, cancer_cell = TumorCellType,
                        cluster_type = opt$cluster_type)

run_cnv(obj = obj, thread = opt$thread)